#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
import cv2
import numpy as np
from cv_bridge import CvBridge
from sensor_msgs.msg import Image, CameraInfo
from darknet_ros_msgs.msg import BoundingBoxes

class CameraModule:
    def __init__(self):
        rospy.init_node('camera_module')
        rospy.loginfo("初始化摄像头模块")

        self.bridge = CvBridge()
        self.color_image = None
        self.depth_image = None
        self.color_camera_info = None
        self.depth_camera_info = None

        # 订阅bounding_boxes、RGB图像、深度图像和相机信息话题
        self.bounding_boxes_subscriber = rospy.Subscriber('/darknet_ros/bounding_boxes', BoundingBoxes, self.bounding_boxes_callback)
        self.color_image_subscriber = rospy.Subscriber('/camera/color/image_raw', Image, self.color_image_callback)
        self.depth_image_subscriber = rospy.Subscriber('/camera/depth/image_raw', Image, self.depth_image_callback)
        self.color_camera_info_subscriber = rospy.Subscriber('/camera/color/camera_info', CameraInfo, self.color_camera_info_callback)
        self.depth_camera_info_subscriber = rospy.Subscriber('/camera/depth/camera_info', CameraInfo, self.depth_camera_info_callback)

        rospy.loginfo("订阅bounding_boxes、RGB图像、深度图像和相机信息话题")

    def color_image_callback(self, msg):
        self.color_image = self.bridge.imgmsg_to_cv2(msg, 'bgr8')
        rospy.loginfo("接收到RGB图像，分辨率: %dx%d", self.color_image.shape[1], self.color_image.shape[0])

    def depth_image_callback(self, msg):
        self.depth_image = self.bridge.imgmsg_to_cv2(msg, '32FC1')
        rospy.loginfo("接收到深度图像，分辨率: %dx%d", self.depth_image.shape[1], self.depth_image.shape[0])

    def color_camera_info_callback(self, msg):
        self.color_camera_info = msg

    def depth_camera_info_callback(self, msg):
        self.depth_camera_info = msg

    def bounding_boxes_callback(self, msg):
        rospy.loginfo("接收bounding_boxes话题消息")

        if self.color_image is None or self.depth_image is None or self.depth_camera_info is None:
            rospy.logwarn("尚未接收到完整的图像和相机信息")
            return

        # 确保深度图像和颜色图像的分辨率一致
        if self.color_image.shape[:2] != self.depth_image.shape[:2]:
            depth_image_resized = cv2.resize(self.depth_image, (self.color_image.shape[1], self.color_image.shape[0]), interpolation=cv2.INTER_NEAREST)
        else:
            depth_image_resized = self.depth_image

        for box in msg.bounding_boxes:
            if box.Class == 'apple':  # 只处理类别为"apple"的目标
                rospy.loginfo("目标类别: %s, Bounding Box: (%d, %d, %d, %d)", box.Class, box.xmin, box.ymin, box.xmax, box.ymax)
                
                # 计算目标中心点的像素坐标
                pixel_x = int((box.xmin + box.xmax) / 2)
                pixel_y = int((box.ymin + box.ymax) / 2)
                
                # 确保像素坐标在深度图像的范围内
                if pixel_x < 0 or pixel_x >= depth_image_resized.shape[1] or pixel_y < 0 or pixel_y >= depth_image_resized.shape[0]:
                    rospy.logwarn("像素坐标超出深度图像范围")
                    continue
                
                # 获取目标中心点的深度值
                depth = depth_image_resized[pixel_y, pixel_x]
                
                # 将像素坐标转换为相机坐标
                fx = self.depth_camera_info.K[0]
                fy = self.depth_camera_info.K[4]
                cx = self.depth_camera_info.K[2]
                cy = self.depth_camera_info.K[5]
                
                x = (pixel_x - cx) * depth / fx
                y = (pixel_y - cy) * depth / fy
                z = depth
                
                rospy.loginfo("目标在相机坐标系中的位置: (%.2f, %.2f, %.2f)", x, y, z)

if __name__ == '__main__':
    try:
        camera_module = CameraModule()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
